package com.varian.auth.service.impl;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.collection.ListUtil;
import cn.hutool.core.convert.Convert;
import cn.hutool.core.date.DateUtil;
import cn.hutool.core.date.LocalDateTimeUtil;
import cn.hutool.core.lang.Assert;
import cn.hutool.core.util.StrUtil;
import com.baomidou.mybatisplus.extension.conditions.query.LambdaQueryChainWrapper;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.Module;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.varian.auth.constant.AuthConstant;
import com.varian.auth.convert.OAuth2AuthorizationInfoConvert;
import com.varian.auth.mapper.OAuth2AuthorizationInfoMapper;
import com.varian.auth.model.dto.OAuth2AuthorizationInfoPageDTO;
import com.varian.auth.model.entity.OAuth2AuthorizationInfo;
import com.varian.auth.model.vo.OAuth2AuthorizationInfoItemVO;
import com.varian.auth.service.IOAuth2AuthorizationInfoService;
import com.varian.auth.service.IOAuth2RegisteredClientService;
import com.varian.mybatis.base.BaseServiceImpl;
import com.varian.mybatis.util.PageUtil;
import com.varian.tool.exception.BizException;
import org.springframework.dao.DataRetrievalFailureException;
import org.springframework.security.jackson2.SecurityJackson2Modules;
import org.springframework.security.oauth2.core.*;
import org.springframework.security.oauth2.core.endpoint.OAuth2ParameterNames;
import org.springframework.security.oauth2.core.oidc.OidcIdToken;
import org.springframework.security.oauth2.core.oidc.endpoint.OidcParameterNames;
import org.springframework.security.oauth2.server.authorization.JdbcOAuth2AuthorizationService;
import org.springframework.security.oauth2.server.authorization.OAuth2Authorization;
import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationCode;
import org.springframework.security.oauth2.server.authorization.OAuth2TokenType;
import org.springframework.security.oauth2.server.authorization.client.RegisteredClient;
import org.springframework.security.oauth2.server.authorization.jackson2.OAuth2AuthorizationServerJackson2Module;
import org.springframework.stereotype.Service;

import java.time.Instant;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * @author ben
 * @since 2024/6/30
 */
@Service
public class OAuth2AuthorizationInfoInfoServiceImpl extends BaseServiceImpl<OAuth2AuthorizationInfoMapper, OAuth2AuthorizationInfo> implements IOAuth2AuthorizationInfoService {

    private final OAuth2AuthorizationInfoConvert baseConvert;
    private final IOAuth2RegisteredClientService auth2RegisteredClientService;
    private final ObjectMapper objectMapper;

    public OAuth2AuthorizationInfoInfoServiceImpl(OAuth2AuthorizationInfoConvert baseConvert, IOAuth2RegisteredClientService auth2RegisteredClientService) {
        this.baseConvert = baseConvert;
        this.auth2RegisteredClientService = auth2RegisteredClientService;

        this.objectMapper = new ObjectMapper();
        ClassLoader classLoader = JdbcOAuth2AuthorizationService.class.getClassLoader();
        List<Module> securityModules = SecurityJackson2Modules.getModules(classLoader);
        this.objectMapper.registerModules(securityModules);
        this.objectMapper.registerModule(new OAuth2AuthorizationServerJackson2Module());
    }

    @Override
    public void save(OAuth2Authorization authorization) {
        Assert.notNull(authorization, () -> BizException.of(AuthConstant.AUTHORIZATION_CANNOT_BE_EMPTY));
        OAuth2AuthorizationInfo existingAuthorization = getById(authorization.getId());
        OAuth2AuthorizationInfo authorizationInfo = convert(authorization);
        if (existingAuthorization == null) {
            save(authorizationInfo);
        } else {
            updateById(authorizationInfo);
        }
    }

    @Override
    public void remove(OAuth2Authorization authorization) {
        Assert.notNull(authorization, () -> BizException.of(AuthConstant.AUTHORIZATION_CANNOT_BE_EMPTY));
        removeById(authorization.getId());
    }

    @Override
    public OAuth2Authorization findById(String id) {
        return findBy(lambdaQuery().eq(OAuth2AuthorizationInfo::getId, id));
    }

    @Override
    public OAuth2Authorization findByToken(String token, OAuth2TokenType tokenType) {
        Assert.notEmpty(token, () -> BizException.of(AuthConstant.TOKEN_CANNOT_BE_EMPTY));
        LambdaQueryChainWrapper<OAuth2AuthorizationInfo> wrapper = lambdaQuery();
        switch (token) {
            case OAuth2ParameterNames.STATE -> wrapper.eq(OAuth2AuthorizationInfo::getState, token);
            case OAuth2ParameterNames.CODE -> wrapper.eq(OAuth2AuthorizationInfo::getAuthorizationCodeValue, token);
            case OAuth2ParameterNames.ACCESS_TOKEN -> wrapper.eq(OAuth2AuthorizationInfo::getAccessTokenValue, token);
            case OidcParameterNames.ID_TOKEN -> wrapper.eq(OAuth2AuthorizationInfo::getOidcIdTokenValue, token);
            case OAuth2ParameterNames.REFRESH_TOKEN -> wrapper.eq(OAuth2AuthorizationInfo::getRefreshTokenValue, token);
            case OAuth2ParameterNames.USER_CODE -> wrapper.eq(OAuth2AuthorizationInfo::getUserCodeValue, token);
            case OAuth2ParameterNames.DEVICE_CODE -> wrapper.eq(OAuth2AuthorizationInfo::getDeviceCodeValue, token);
            default -> wrapper.eq(OAuth2AuthorizationInfo::getAuthorizationCodeValue, token).or()
                    .eq(OAuth2AuthorizationInfo::getAccessTokenValue, token).or()
                    .eq(OAuth2AuthorizationInfo::getOidcIdTokenValue, token).or()
                    .eq(OAuth2AuthorizationInfo::getRefreshTokenValue, token).or()
                    .eq(OAuth2AuthorizationInfo::getUserCodeValue, token).or()
                    .eq(OAuth2AuthorizationInfo::getDeviceCodeValue, token);
        }
        return findBy(wrapper);
    }

    @Override
    public Page<OAuth2AuthorizationInfoItemVO> page(OAuth2AuthorizationInfoPageDTO dto) {
        Page<OAuth2AuthorizationInfo> page = baseMapper.selectPage(dto.toPage(), dto);
        return PageUtil.convert(page, baseConvert::toItemVO);
    }

    private OAuth2Authorization findBy(LambdaQueryChainWrapper<OAuth2AuthorizationInfo> wrapper) {
        OAuth2AuthorizationInfo auth2AuthorizationInfo = baseMapper.selectOne(wrapper);
        if (auth2AuthorizationInfo == null) {
            return null;
        }
        return convert(auth2AuthorizationInfo);
    }


    private OAuth2AuthorizationInfo convert(OAuth2Authorization authorization) {
        OAuth2AuthorizationInfo info = new OAuth2AuthorizationInfo();
        info.setRegisteredClientId(Convert.toLong(authorization.getRegisteredClientId()))
                .setPrincipalName(authorization.getPrincipalName())
                .setAuthorizationGrantType(authorization.getAuthorizationGrantType().getValue())
                .setAuthorizedScopes(ListUtil.list(false, authorization.getAuthorizedScopes()))
                .setAttributes(writeMap(authorization.getAttributes()))
                .setState(authorization.getAttribute(OAuth2ParameterNames.STATE));
        OAuth2Authorization.Token<OAuth2AuthorizationCode> authorizationCode = authorization.getToken(OAuth2AuthorizationCode.class);
        if (authorizationCode != null) {
            info.setAuthorizationCodeValue(authorizationCode.getToken().getTokenValue())
                    .setAuthorizationCodeIssuedAt(LocalDateTimeUtil.of(authorizationCode.getToken().getIssuedAt()))
                    .setAuthorizationCodeExpiresAt(LocalDateTimeUtil.of(authorizationCode.getToken().getExpiresAt()))
                    .setAuthorizationCodeMetadata(writeMap(authorizationCode.getMetadata()));
        }
        OAuth2Authorization.Token<OAuth2AccessToken> accessToken = authorization.getToken(OAuth2AccessToken.class);
        if (accessToken != null) {
            info.setAccessTokenValue(accessToken.getToken().getTokenValue())
                    .setAccessTokenType(accessToken.getToken().getTokenType().getValue())
                    .setAccessTokenScopes(ListUtil.list(false, accessToken.getToken().getScopes()))
                    .setAccessTokenIssuedAt(LocalDateTimeUtil.of(accessToken.getToken().getIssuedAt()))
                    .setAccessTokenExpiresAt(LocalDateTimeUtil.of(accessToken.getToken().getExpiresAt()))
                    .setAccessTokenMetadata(writeMap(accessToken.getMetadata()));
        }
        OAuth2Authorization.Token<OidcIdToken> oidcIdToken = authorization.getToken(OidcIdToken.class);
        if (oidcIdToken != null) {
            info.setOidcIdTokenValue(oidcIdToken.getToken().getTokenValue())
                    .setOidcIdTokenIssuedAt(LocalDateTimeUtil.of(oidcIdToken.getToken().getIssuedAt()))
                    .setOidcIdTokenExpiresAt(LocalDateTimeUtil.of(oidcIdToken.getToken().getExpiresAt()))
                    .setOidcIdTokenMetadata(writeMap(oidcIdToken.getMetadata()));
        }
        OAuth2Authorization.Token<OAuth2RefreshToken> refreshToken = authorization.getRefreshToken();
        if (refreshToken != null) {
            info.setRefreshTokenValue(refreshToken.getToken().getTokenValue())
                    .setRefreshTokenIssuedAt(LocalDateTimeUtil.of(refreshToken.getToken().getIssuedAt()))
                    .setRefreshTokenExpiresAt(LocalDateTimeUtil.of(refreshToken.getToken().getExpiresAt()))
                    .setRefreshTokenMetadata(writeMap(refreshToken.getMetadata()));
        }
        OAuth2Authorization.Token<OAuth2UserCode> userCode = authorization.getToken(OAuth2UserCode.class);
        if (userCode != null) {
            info.setUserCodeValue(userCode.getToken().getTokenValue())
                    .setUserCodeIssuedAt(LocalDateTimeUtil.of(userCode.getToken().getIssuedAt()))
                    .setUserCodeExpiresAt(LocalDateTimeUtil.of(userCode.getToken().getExpiresAt()))
                    .setUserCodeMetadata(writeMap(userCode.getMetadata()));
        }
        OAuth2Authorization.Token<OAuth2DeviceCode> deviceCode = authorization.getToken(OAuth2DeviceCode.class);
        if (deviceCode != null) {
            info.setDeviceCodeValue(deviceCode.getToken().getTokenValue())
                    .setDeviceCodeIssuedAt(LocalDateTimeUtil.of(deviceCode.getToken().getIssuedAt()))
                    .setDeviceCodeExpiresAt(LocalDateTimeUtil.of(deviceCode.getToken().getExpiresAt()))
                    .setDeviceCodeMetadata(writeMap(deviceCode.getMetadata()));
        }
        return info;
    }


    @SuppressWarnings("unchecked")
    private OAuth2Authorization convert(OAuth2AuthorizationInfo auth2AuthorizationInfo) {
        if (auth2AuthorizationInfo == null) {
            return null;
        }
        Long registeredClientId = auth2AuthorizationInfo.getRegisteredClientId();
        RegisteredClient registeredClient = auth2RegisteredClientService.findById(registeredClientId);
        Assert.notNull(registeredClient, () -> new DataRetrievalFailureException("The RegisteredClient with id '" + registeredClientId + "' was not found in the RegisteredClientRepository."));
        OAuth2Authorization.Builder builder = OAuth2Authorization.withRegisteredClient(registeredClient);
        Set<String> authorizedScopes = Collections.emptySet();
        if (CollUtil.isNotEmpty(auth2AuthorizationInfo.getAuthorizedScopes())) {
            authorizedScopes = CollUtil.newHashSet(auth2AuthorizationInfo.getAuthorizedScopes());
        }
        builder.id(auth2AuthorizationInfo.getId().toString())
                .principalName(auth2AuthorizationInfo.getPrincipalName())
                .authorizationGrantType(new AuthorizationGrantType(auth2AuthorizationInfo.getAuthorizationGrantType()))
                .authorizedScopes(authorizedScopes)
                .attributes(attrs -> attrs.putAll(parseMap(auth2AuthorizationInfo.getAttributes())));
        if (StrUtil.isNotEmpty(auth2AuthorizationInfo.getState())) {
            builder.attribute(OAuth2ParameterNames.STATE, auth2AuthorizationInfo.getState());
        }
        Instant tokenIssuedAt;
        Instant tokenExpiresAt;
        String authorizationCodeValue = auth2AuthorizationInfo.getAuthorizationCodeValue();
        if (StrUtil.isNotEmpty(authorizationCodeValue)) {
            tokenIssuedAt = DateUtil.toInstant(auth2AuthorizationInfo.getAuthorizationCodeIssuedAt());
            tokenExpiresAt = DateUtil.toInstant(auth2AuthorizationInfo.getAuthorizationCodeExpiresAt());
            OAuth2AuthorizationCode auth2AuthorizationCode = new OAuth2AuthorizationCode(authorizationCodeValue, tokenIssuedAt, tokenExpiresAt);
            builder.token(auth2AuthorizationCode, metadata -> metadata.putAll(parseMap(auth2AuthorizationInfo.getAuthorizationCodeMetadata())));
        }
        String accessTokenValue = auth2AuthorizationInfo.getAccessTokenValue();
        if (StrUtil.isNotEmpty(accessTokenValue)) {
            tokenIssuedAt = DateUtil.toInstant(auth2AuthorizationInfo.getAccessTokenIssuedAt());
            tokenExpiresAt = DateUtil.toInstant(auth2AuthorizationInfo.getAccessTokenExpiresAt());
            OAuth2AccessToken.TokenType tokenType = null;
            if (OAuth2AccessToken.TokenType.BEARER.getValue().equalsIgnoreCase(auth2AuthorizationInfo.getAccessTokenType())) {
                tokenType = OAuth2AccessToken.TokenType.BEARER;
            }
            Set<String> scopes = Collections.emptySet();
            if (CollUtil.isNotEmpty(auth2AuthorizationInfo.getAccessTokenScopes())) {
                scopes = CollUtil.newHashSet(auth2AuthorizationInfo.getAccessTokenScopes());
            }
            OAuth2AccessToken auth2AccessToken = new OAuth2AccessToken(tokenType, accessTokenValue, tokenIssuedAt, tokenExpiresAt, scopes);
            builder.token(auth2AccessToken, metadata -> metadata.putAll(parseMap(auth2AuthorizationInfo.getAccessTokenMetadata())));
        }
        String oidcIdTokenValue = auth2AuthorizationInfo.getOidcIdTokenValue();
        if (StrUtil.isNotEmpty(oidcIdTokenValue)) {
            tokenIssuedAt = DateUtil.toInstant(auth2AuthorizationInfo.getOidcIdTokenIssuedAt());
            tokenExpiresAt = DateUtil.toInstant(auth2AuthorizationInfo.getOidcIdTokenExpiresAt());
            Map<String, Object> oidcTokenMetadata = parseMap(auth2AuthorizationInfo.getOidcIdTokenMetadata());
            OidcIdToken oidcToken = new OidcIdToken(oidcIdTokenValue, tokenIssuedAt, tokenExpiresAt,
                    (Map<String, Object>) oidcTokenMetadata.get(OAuth2Authorization.Token.CLAIMS_METADATA_NAME));
            builder.token(oidcToken, (metadata) -> metadata.putAll(oidcTokenMetadata));
        }
        String refreshTokenValue = auth2AuthorizationInfo.getRefreshTokenValue();
        if (StrUtil.isNotEmpty(refreshTokenValue)) {
            tokenIssuedAt = DateUtil.toInstant(auth2AuthorizationInfo.getRefreshTokenIssuedAt());
            tokenExpiresAt = DateUtil.toInstant(auth2AuthorizationInfo.getRefreshTokenExpiresAt());
            OAuth2RefreshToken auth2RefreshToken = new OAuth2RefreshToken(refreshTokenValue, tokenIssuedAt, tokenExpiresAt);
            builder.token(auth2RefreshToken, metadata -> metadata.putAll(parseMap(auth2AuthorizationInfo.getRefreshTokenMetadata())));
        }
        String userCodeValue = auth2AuthorizationInfo.getUserCodeValue();
        if (StrUtil.isNotEmpty(userCodeValue)) {
            tokenIssuedAt = DateUtil.toInstant(auth2AuthorizationInfo.getUserCodeIssuedAt());
            tokenExpiresAt = DateUtil.toInstant(auth2AuthorizationInfo.getUserCodeExpiresAt());
            OAuth2UserCode userCode = new OAuth2UserCode(userCodeValue, tokenIssuedAt, tokenExpiresAt);
            builder.token(userCode, metadata -> metadata.putAll(parseMap(auth2AuthorizationInfo.getUserCodeMetadata())));
        }
        String deviceCodeValue = auth2AuthorizationInfo.getDeviceCodeValue();
        if (StrUtil.isNotEmpty(deviceCodeValue)) {
            tokenIssuedAt = DateUtil.toInstant(auth2AuthorizationInfo.getDeviceCodeIssuedAt());
            tokenExpiresAt = DateUtil.toInstant(auth2AuthorizationInfo.getDeviceCodeExpiresAt());
            OAuth2DeviceCode deviceCode = new OAuth2DeviceCode(deviceCodeValue, tokenIssuedAt, tokenExpiresAt);
            builder.token(deviceCode, metadata -> metadata.putAll(parseMap(auth2AuthorizationInfo.getDeviceCodeMetadata())));
        }
        return builder.build();
    }

    private Map<String, Object> parseMap(String data) {
        try {
            return this.objectMapper.readValue(data, new TypeReference<>() {
            });
        } catch (Exception ex) {
            throw new IllegalArgumentException(ex.getMessage(), ex);
        }
    }

    private String writeMap(Map<String, Object> data) {
        try {
            return this.objectMapper.writeValueAsString(data);
        } catch (Exception ex) {
            throw new IllegalArgumentException(ex.getMessage(), ex);
        }
    }
}
